//
// Created by mio on 25-4-17.
//

#include "RtmpChunk.h"

#include <cstring>

int RtmpChunk::streamid_ = 0;

RtmpChunk::RtmpChunk()
{
    state_ = PARSE_HEADER;
    chunk_stream_id_ = -1;
    ++streamid_;
}

RtmpChunk::~RtmpChunk()
{

}

int RtmpChunk::Parse(BufferReader& in_buffer, RtmpMessage& out_rtmp_msg)
{
    int ret = 0;
    if (!in_buffer.ReadableBytes()) return 0;
    if (state_ == PARSE_HEADER) ret = ParseChunkHeader(in_buffer);
    else
    {
        ret = ParseChunkBody(in_buffer);
        // 如果 Body 解析成功 (ret > 0) 并且我们知道当前的 CSID (chunl_stream_id_ >= 0)
        if (ret > 0 && chunk_stream_id_ >= 0) // 解析成功
        {
            // 获取当前 CSID 对应的正在组装的消息对象引用
            auto& rtmp_msg = rtmp_message_[chunk_stream_id_];
            if (rtmp_msg.index == rtmp_msg.length) // 说明message是一个完整的
            {
                // 处理绝对 / 扩展时间戳
                if (rtmp_msg.timestamp >= 0xffffff)
                {
                    // 使用 32 位扩展时间戳作为绝对时间戳 (通常在 fmt=0 时)
                    // 或者应用扩展时间戳（具体逻辑依赖 RTMP 规范细节和 fmt 类型，这里直接赋值给 _timestamp 可能是简化处理）
                    // 根据 RTMP Spec，如果扩展时间戳存在，它 *总是* 绝对时间戳。
                    // 所以这里直接用 extend_timestamp 覆盖/设置 _timestamp 是合理的。
                    rtmp_msg._timestamp += rtmp_msg.extend_timestamp;
                }
                else
                {
                    rtmp_msg._timestamp += rtmp_msg.timestamp;
                }

                out_rtmp_msg = rtmp_msg;    // 输出到调用方
                chunk_stream_id_ = -1;      // 结束本轮 CSID
                rtmp_msg.Clear();           // 重置缓存
            }
        }
    }
    return ret; // 消费的字节数
}

int RtmpChunk::CreateChunk(uint32_t csid, RtmpMessage& in_msg, char* buf, uint32_t buf_size)
{
    uint32_t buf_offset = 0, payload_offset = 0;
    uint32_t capacity = in_msg.length + in_msg.length / out_chunk_size_ * 5; // 准备的容量（估算）
    if (buf_size < capacity) // 缓冲区大小小于准备的容量的话就返回
    {
        return -1;
    }
    buf_offset += CreateBasicHeader(0, csid, buf + buf_offset);
    buf_offset += CreateMessageHeader(0, in_msg, buf + buf_offset);
    if (in_msg._timestamp >= 0xffffff)
    {
        WriteUint32BE((char*)buf + buf_offset,(uint32_t)in_msg.extend_timestamp);
        buf_offset += 4;
    }
    // --- 循环写入 Chunk Payload 和后续 Chunk Header ---
    while (in_msg.length > 0)
    {
        if (in_msg.length > out_chunk_size_)
        {
            memcpy(buf + buf_offset, in_msg.payload.get() + payload_offset, out_chunk_size_);
            payload_offset += out_chunk_size_;
            buf_offset += out_chunk_size_;
            in_msg.length -= out_chunk_size_;

            buf_offset += CreateBasicHeader(3, csid, buf + buf_offset);
            if (in_msg._timestamp >= 0xffffff)
            {
                WriteUint32BE(buf + buf_offset, (uint32_t)in_msg.extend_timestamp);
                buf_offset += 4;
            }
        }
        else // 最后一个包
        {
            memcpy(buf + buf_offset, in_msg.payload.get() + payload_offset, in_msg.length);
            buf_offset += in_msg.length;
            in_msg.length = 0;
            break;
        }
    }
    return buf_offset;
}

int RtmpChunk::ParseChunkHeader(BufferReader& buffer)
{
    uint32_t bytes_used = 0;
    uint8_t* buf = (uint8_t*) buffer.Peek();
    uint32_t buf_size = buffer.ReadableBytes();

    /* —— Basic Header —— */
    uint8_t flags = buf[bytes_used];
    uint8_t fmt = (flags >> 6);
    if (fmt >= 4) // fmt 0-3
    {
        return -1;
    }

    bytes_used += 1;
    uint8_t csid = flags & 0x3f;
    if (csid == 0) // 两字节 CSID
    {
        if (buf_size < (bytes_used + 2))
        {
            return 0;
        }
        csid += buf[bytes_used] + 64;
        bytes_used += 1;
    }
    else if (csid == 1) // 三字节 CSID
    {
        if (buf_size < (bytes_used + 3))
        {
            return 0;
        }
        csid += buf[bytes_used + 1] * 255 + buf[bytes_used] + 64;
        bytes_used += 2;
    }
    /* —— Message Header 长度 —— */
    uint32_t header_len = KChunkMessageHeaderLength[fmt];
    if (buf_size < (header_len + bytes_used)) return 0;

    RtmpMessageHeader header;
    memset(&header, 0, sizeof(RtmpMessageHeader));
    memcpy(&header, buf + bytes_used, header_len);
    bytes_used += header_len;

    auto& rtmp_msg = rtmp_message_[csid];
    chunk_stream_id_ = rtmp_msg.csid = csid;
    if (fmt == 0 || fmt == 1)
    {
        uint32_t length = ReadUint24BE((char*)header.length);
        if (rtmp_msg.length != length || !rtmp_msg.payload)
        {
            rtmp_msg.length = length;
            rtmp_msg.payload.reset(new char[rtmp_msg.length], std::default_delete<char[]>());
        }
        rtmp_msg.index = 0;
        rtmp_msg.type_id = header.type_id;
    }

    // --- 更新 Message Length, Type ID ---
    if (fmt == 0)
    {
        rtmp_msg.stream_id = ReadUint32LE((char*)header.stream_id);
    }
    uint32_t timestamp = ReadUint24BE((char*)header.timestamp);
    uint32_t extend_timestamp = 0;
    if (timestamp >= 0xffffff || rtmp_msg.timestamp >= 0xffffff)
    {
        if (buf_size < (bytes_used + 4)) return 0;
        extend_timestamp = ReadUint32BE((char*)buf + bytes_used);
        bytes_used += 4;
    }
    // --- 计算和应用时间戳 ---
    // 这部分是 RTMP 时间戳处理的核心
    // rtmp_msg._timestamp 存储计算后的绝对时间戳
    // rtmp_msg.timestamp 存储本次从 Header 读取的 24 位值 (可能是绝对值或 delta)
    // rtmp_msg.extend_timestamp 存储本次从 Header 读取的 32 位值 (绝对值或 delta)

    // 何时更新时间戳？当一个新的消息开始时（fmt=0），或者一个消息的后续块带有时间戳增量时（fmt=1, fmt=2）


    // rtmp_msg.index == 0 表示这是消息的第一个块（或者刚刚被重置）
    if (rtmp_msg.index == 0)
    {
        if (fmt == 0)
        {
            rtmp_msg._timestamp = 0;
            rtmp_msg.timestamp = timestamp;
            rtmp_msg.extend_timestamp = extend_timestamp;
        }
        else
        {
            if (rtmp_msg.timestamp >= 0xffffff)
            {
                rtmp_msg.extend_timestamp += extend_timestamp;
            }
            else
            {
                rtmp_msg.timestamp += timestamp;
            }
        }
    }

    state_ = PARSE_BODY;
    buffer.Retrieve(bytes_used);
    return bytes_used;
}

int RtmpChunk::ParseChunkBody(BufferReader& buffer)
{
    uint32_t bytes_used = 0;
    uint8_t* buf = (uint8_t*)buffer.Peek();
    uint32_t buf_size = buffer.ReadableBytes();

    if (chunk_stream_id_ < 0) return -1;

    auto& rtmp_msg = rtmp_message_[chunk_stream_id_];
    // 计算本次需要读取的 Payload 大小 (chunk_size)
    // 它是以下两者中的较小值：
    // 1. 整个消息剩余未读取的长度 (rtmp_msg.length - rtmp_msg.index)
    // 2. RTMP 协议规定的当前 Chunk 最大大小 (in_chunk_size_)
    uint32_t chunk_size = rtmp_msg.length - rtmp_msg.index;
    if (chunk_size > in_chunk_size_) chunk_size = in_chunk_size_;
    // 检查缓冲区是否有足够的数据来读取这个 chunk_size
    if (buf_size < (chunk_size + bytes_used)) return 0;
    // 检查是否有逻辑错误（例如 index 超出 length）
    if (rtmp_msg.index + chunk_size > rtmp_msg.length) return -1;
    // 从缓冲区拷贝数据到 RtmpMessage 的 payload 中正确的位置
    memcpy(rtmp_msg.payload.get() + rtmp_msg.index, buf + bytes_used, chunk_size);
    bytes_used += chunk_size;
    rtmp_msg.index += chunk_size;
    // 检查当前 Chunk 是否接收完毕
    // 接收完毕的条件是：
    // 1. 整个消息已经接收完整 (rtmp_msg.index == rtmp_msg.length)
    // 2. 或者，当前 Chunk 已经达到了协商的 Chunk 大小 (rtmp_msg.index % in_chunk_size_ == 0)
    //    注意：这个条件对于最后一个不足 in_chunk_size_ 的 chunk 也成立，因为它同时满足条件1。
    //    所以主要判断是：是否读完了一个完整的 `in_chunk_size_` 块，或者整个消息读完了。
    if (rtmp_msg.index >= rtmp_msg.length || rtmp_msg.index % in_chunk_size_ == 0)
    {
        state_ = PARSE_HEADER;
    }
    buffer.Retrieve(bytes_used);
    return bytes_used;
}

int RtmpChunk::CreateBasicHeader(uint8_t fmt, uint32_t csid, char* buf)
{
    int len = 0;
    // CSID 范围: 2 到 65599
    // 格式:
    // 1 byte: fmt(2b) + csid(6b)  (csid 2-63)
    // 2 byte: fmt(2b) + 0(6b), csid-64 (1B) (csid 64-319)
    // 3 byte: fmt(2b) + 1(6b), csid-64 (2B, LittleEndian) (csid 320-65599)

    if (csid > 64 + 255) // 说明这个basic头占3字节
    {
        buf[len++] = (fmt << 6) | 1;
        // 写入 2 字节的小端数据
        buf[len++] = (csid - 64) & 0xff;
        buf[len++] = ((csid - 64) >> 8) & 0xff;
    }
    else if (csid >= 64) // 占3字节
    {
        buf[len++] = (fmt << 6) | 0;
        buf[len++] = (csid - 64) & 0xff;
    }
    else // 1字节
    {
        buf[len++] = (fmt << 6) | csid;
    }
    return len;
}

int RtmpChunk::CreateMessageHeader(uint8_t fmt, RtmpMessage& rtmp_msg, char* buf)
{
    int len = 0;
    // --- 写入 Timestamp (或 Delta) ---
    if (fmt <= 2)
    {
        if (rtmp_msg.timestamp < 0xffffff)
        {
            WriteUint24BE((char*)buf, (uint32_t)rtmp_msg._timestamp);
        }
        else
        {
            WriteUint24BE((char*)buf, 0xffffff);
        }
        len += 3;
    }
    // --- 写入 Message Length 和 Type ID ---
    if (fmt <= 1)
    {
        WriteUint24BE((char*)buf + len, rtmp_msg.length);
        len += 3;
        buf[len++] = rtmp_msg.type_id;
    }
    // --- 写入 Message Stream ID ---
    if (fmt == 0)
    {
        WriteUint32LE((char*)buf + len, rtmp_msg.stream_id);
        len += 4;
    }
    return len;
}
